package comlib.chisel

import chisel3._

/*
 Define counter class
 */

class Counter[T<:Bits with Num[T]] (init:T) {

    object Step extends Enumeration {
        type Step = Value
        val init, load, boundary, count, wrap, output = Value
    }

    protected val zeroValue = 0.U.asInstanceOf[T]
    protected val allOneValue = Seq.fill(init.getWidth)("1").mkString("b","","").U.asInstanceOf[T]
    protected val oneValue = 1.U.asInstanceOf[T]

    protected var mode = "LOOP2"
    protected var step = Step.load

    protected val _value_ = WireInit(init)      // Inner counter value
    protected val _value_next_ = WireInit(init)
    protected val _overflow_ = WireInit(false.B)
    protected val _underflow_ = WireInit(false.B)

    /* Set max ans min signals */
    protected var _max_ = allOneValue
    protected var _min_ = zeroValue

    def boundary( max:T, min:T=0.U.asInstanceOf[T], mode:String="LOOP" ) = {
        _max_ = max
        _min_ = min
        this.mode = mode
        this
    }
    def max = _max_
    def min = _min_

    /* Interupt */
    protected val _overscope_ = WireInit(false.B)
    def overscope = _overscope_

    private var isLoaded = false

    // ---- Wrap
    // private val _wrap_ = WireInit(false.B)

    // /** Return true.B when the counter wrap from 0xF..F to zero or from zero to 0xF..F, when have not set max and min boundary
    //   * 
    //   * @see [[boundary]]
    //   * @param lock Register lock output flag
    //   */
    // def wrap( lock: Boolean=false ) = LockerInit( lock, false.B, _wrap_ )


    //-----------  New version architecture  -------------
    private var incFn : Option[()=>T] = None
    private var decFn : Option[()=>T] = None
 
    private def autoStepToOutput() = {
        if( step < Step.count ) autoStepCount() // increase()
        else require( step < Step.output, "Counter had a `OUTPUT` call" )    
    }

    /* Get counter value */
    def value( isLock:Boolean = true ) : T = {
        autoStepToOutput()

        val valLgc = if( isLock ) RegInit( init ) else WireInit( init )     // Logic circuit output

        valLgc := _value_next_   // Transmit by combinatiion or sequential circuit
        if( isLock )
            _value_ := valLgc

        valLgc
    }

    /* Get counter value and overflow flag */
    def value( isLock:Boolean, ofLock:Boolean, ufLock:Boolean ) = {
        autoStepToOutput()

        val valLgc = if( isLock ) RegInit( init ) else WireInit( init )     // Logic circuit output
        val ofLgc = if( ofLock ) RegInit( false.B ) else WireInit( false.B )
        val ufLgc = if( ufLock ) RegInit( false.B ) else WireInit( false.B )

        valLgc := _value_next_
        ofLgc := _overflow_
        ufLgc := _underflow_

        if( isLock )
            _value_ := valLgc

        (valLgc, ofLgc, ufLgc)
    }

    def load(value: T) = {

        require( step  == Step.load, "load step error")

        _value_ := value
        isLoaded = true
        this
    }

    private def autoStepCount() = {
        if( incFn != None || decFn != None )
            count()
        else
            count( true.B, oneValue, false.B, zeroValue )
    }
    
    private def commonWrapProc() = {
        if( step < Step.wrap ) autoStepCount()
            // increase(true.B)
        step = Step.wrap
    }

    /* Reset counter value to init */
    def forceInit(rst: Bool) = {
        commonWrapProc()
        when( rst ){ _value_next_ := init }
        this
    }

    /* Clear counter: mean set counter value to '0' */
    def force0(clr: Bool) = {
        commonWrapProc()
        when( clr ){ _value_next_ := zeroValue }
        this
    }

    /* Set counter value to XXX */
    def force( en: Bool, value : T ) = {
        commonWrapProc()
        when( en ){ _value_next_ := value }
        this
    }

    def forceMax( en: Bool ) = {
        commonWrapProc()
        when( en ){ _value_next_ := _max_ }
        this
    }

    def forceMin( en: Bool ) = {
        commonWrapProc()
        when( en ){ _value_next_ := min }
        this
    }

    def force1( en: Bool ) = {
        commonWrapProc()
        when( en ){ _value_next_ := allOneValue }
        this
    }

    

    /* Main count process */
    def count( inc: Bool, ival: T, dec: Bool, dval: T ) = {

        require( step <= Step.count, "Counter cannot work in `WRAP` station")

        _overflow_ := false.B
        _underflow_ := false.B
        _overscope_ := false.B
        // _wrap_ := false.B

        when( inc && dec ){
            when( ival >= dval ){

                val diff = ival - dval

                mode match {
                    case "LOOP2" => 
                        _value_next_ := _value_ + diff
                    case "LOOP" => 
                        when( _max_ - _min_ < diff ){
                            _overflow_ := true.B
                            _overscope_ := true.B
                            _value_next_ := _value_      
                        }.elsewhen( _value_ > _max_ - diff ){
                            _overflow_ := true.B
                            _value_next_ := _min_ + (_max_ - diff)
                        }.otherwise{
                            _value_next_ := _value_ + diff
                        } 
                    case _ =>
                        when( _max_ < diff ){
                            _overflow_ := true.B
                            _value_next_ := _max_
                        }.elsewhen( _value_ > _max_ - diff ){
                            _overflow_ := true.B
                            _value_next_ := _max_
                        }.otherwise{
                            _value_next_ := _value_ + diff
                        } 
                }

            }.otherwise{

              val diff = dval - ival
              mode match {
                case "LOOP2" => 
                    _value_next_ := _value_ - diff
                case "LOOP" =>
                  when( _max_ - _min_ < diff ){
                    _underflow_ := true.B
                    _overscope_ := true.B
                    _value_next_ := _value_
                  }.elsewhen( _value_ - _min_ < diff ){      // Overflow
                    _underflow_ := true.B
                    _value_next_ := _min_
                  }.otherwise{
                    _value_next_ := _value_ - (dval-ival)
                  }
                case _ =>
                    when( _value_ - _min_ < diff ){      // Overflow
                        _underflow_ := true.B
                        _value_next_ := _min_
                    }.otherwise{
                        _value_next_ := _value_ - (dval-ival)
                    }
              }
                
            }
        }.elsewhen( inc && !dec ){

            if( mode == "LOOP2" ){
                _value_next_ := _value_ + ival
            }else{
                when( _value_ > _max_ - ival ){
                    _overflow_ := true.B

                    if( mode == "LOOP" ) _value_next_ := _min_ + (_max_ - ival)
                    else _value_next_ := _max_
                }.otherwise{
                    _value_next_ := _value_ + ival
                }
            }

        }.elsewhen( !inc && dec ){

            if( mode == "LOOP2" ){
                _value_next_ := _value_ - dval
            }else{
                when( _value_ < _min_ + dval ){
                    _underflow_ := true.B

                    if( mode == "LOOP" ) _value_next_ := _max_ - (_min_ + dval)
                    else _value_next_ := _min_
                }.otherwise{
                    _value_next_ := _value_ - dval
                }
            }
            
        }.otherwise{
            _value_next_ := _value_
        }

        step = Step.wrap

        this
    }

    def count( inc: Bool, ival: T, dec: Bool ) = {

        _overflow_ := false.B

        when( inc & dec ){
            _value_next_ := ival
        }.elsewhen( inc ){
            if( mode == "LOOP2" ){
                _value_next_ := _value_ + ival
            }else{
                when( _value_ > _max_ - ival ){
                    _overflow_ := true.B

                    if( mode == "LOOP" ) _value_next_ := _min_ + (_max_ - ival)
                    else _value_next_ := _max_
                }.otherwise{
                    _value_next_ := _value_ + ival
                }
            }
        }.elsewhen( dec ){
            _value_next_ := zeroValue
        }.otherwise{
            _value_next_ := _value_
        }

        step = Step.wrap

        this
    }

    def increase( inc:Bool=true.B, ival:T=1.U.asInstanceOf[T] ) = count( inc, ival, false.B, 0.U.asInstanceOf[T] )
    def decrease( dec:Bool=true.B, dval:T=1.U.asInstanceOf[T] ) = count( false.B, 0.U.asInstanceOf[T], dec, dval )

    //-----------  New version ------------
    def inc( inc:Bool, ival:T=1.U.asInstanceOf[T] ) = incFn match {
        case None =>
            this.incFn = Some(() => {
                val iret = WireInit(zeroValue)
                when( inc ){
                    iret := ival
                }
                iret
            })
            this

        case Some(fn) =>
            this.incFn = Some(() => {
                val iret = Wire(chiselTypeOf(init))
                val lret = fn()
                when( inc ){
                    when( ival > allOneValue - lret ){
                        iret := allOneValue
                    }.otherwise{
                        iret := ival + lret
                    }
                }.otherwise{
                    iret := lret
                }
                iret
            })
            this
    }

    def dec( dec:Bool, dval:T=1.U.asInstanceOf[T] ) = decFn match {
        case None =>
            this.decFn = Some(() => {
                val iret = WireInit(zeroValue)
                when( dec ){
                    iret := dval
                }
                iret
            })
            this

        case Some(fn) =>
            this.decFn = Some(() => {
                val iret = Wire(chiselTypeOf(init))
                val lret = fn()
                when( dec ){
                    when( dval > allOneValue - lret ){
                        iret := allOneValue
                    }.otherwise{
                        iret := dval + lret
                    }
                }.otherwise{
                    iret := lret
                }
                iret
            })
            this
    }

    def clear( en:Bool ) = dec( en, _value_ )

    def count(): Counter[T] = (incFn, decFn) match {
        case (None, None) => count( false.B, init, false.B, init )
        case (None, Some(dec)) => count( false.B, init, true.B, dec() )
        case (Some(inc), None) => count( true.B, inc(), false.B, init )
        case (Some(inc), Some(dec)) => count( true.B, inc(), true.B, dec() )
    }
} 

object Counter {

    def apply[T<:Bits with Num[T]]( width:Int=0, init:T=0.U.asInstanceOf[T] ) ={
        if( width == 0 ) new Counter( init )
        else new Counter( 0.U(width.W).asInstanceOf[T] )
    } 

    implicit def CounterToT[T<:Bits with Num[T]](cnt:Counter[T]): T ={
        if( cnt.isLoaded ) cnt.value( false )
        else cnt.value( true )
    }
}
